from sentence_transformers import SentenceTransformer
import numpy as np

def load_model():
    """
    加载预训练的Sentence-BERT模型
    """
    print("正在加载模型...")
    # 使用轻量级模型，适合入门使用
    model = SentenceTransformer('all-MiniLM-L6-v2')
    return model

def encode_sentences(model, sentences):
    """
    将句子转换为向量表示
    """
    print("正在编码句子...")
    # 将句子转换为向量表示
    embeddings = model.encode(sentences)
    return embeddings

def calculate_similarity(embeddings):
    """
    计算句子之间的余弦相似度
    """
    # 计算余弦相似度
    similarity_matrix = np.zeros((len(embeddings), len(embeddings)))
    for i in range(len(embeddings)):
        for j in range(len(embeddings)):
            similarity_matrix[i][j] = np.dot(embeddings[i], embeddings[j]) / (
                np.linalg.norm(embeddings[i]) * np.linalg.norm(embeddings[j])
            )
    return similarity_matrix

def main():
    # 示例句子
    sentences = [
        "我喜欢在公园散步",
        "公园里散步很舒服",
        "今天天气真不错",
        "这个苹果很甜",
        "我最喜欢吃水果了"
    ]
    
    # 加载模型
    model = load_model()
    
    # 获取句子的向量表示
    embeddings = encode_sentences(model, sentences)
    
    # 计算相似度
    similarity_matrix = calculate_similarity(embeddings)
    
    # 打印结果
    print("\n句子相似度矩阵：")
    for i in range(len(sentences)):
        print(f"\n句子 {i+1}: {sentences[i]}")
        print("与其他句子的相似度：")
        for j in range(len(sentences)):
            if i != j:
                print(f"  - 与句子 {j+1} ({sentences[j]}) 的相似度: {similarity_matrix[i][j]:.4f}")

if __name__ == "__main__":
    main()
